/*
 * jngc, Java utility classes by Nicholas Coltharp
 * Copyright (C) 2010 Nicholas Coltharp <ngc1@rice.edu>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package edu.rice.ngc1.util.annotation.processor;

import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Messager;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.Name;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.ArrayType;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.ErrorType;
import javax.lang.model.type.PrimitiveType;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.type.TypeVariable;
import javax.lang.model.util.SimpleElementVisitor6;
import javax.lang.model.util.SimpleTypeVisitor6;
import javax.lang.model.util.Types;
import javax.tools.Diagnostic;

import edu.rice.ngc1.util.annotation.Immutable;
import edu.rice.ngc1.util.annotation.Mutable;
import edu.rice.ngc1.util.annotation.Immutable.Strictness;
import edu.rice.ngc1.util.annotation.processor.tool.Elements;

/**
 * Processes elements marked with the
 * {@link edu.rice.ngc1.util.annotation.Immutable} and
 * {@link edu.rice.ngc1.util.annotation.Mutable} annotations.
 * 
 * For an {@code Immutable} element, the processor first ensures that
 * the element represents a class, and prints an error message if it
 * isn't. It then verifies that the class meets the requirements
 * dictated by the annotation's strictness level, and prints an error
 * message for each problem detected.
 * 
 * For a {@code Mutable} element, the processor ensures that the
 * element belongs to a class marked immutable, and prints an error
 * message if it does not.
 * 
 * Attempting to process an unsupported annotation or element type
 * will result in a an {@link UnsupportedOperationException} being
 * thrown.
 */
@SupportedAnnotationTypes({"edu.rice.ngc1.util.annotation.Immutable",
        "edu.rice.ngc1.util.annotation.Mutable"})
@SupportedSourceVersion(SourceVersion.RELEASE_6)
public final class ImmutableProcessor extends AbstractProcessor {

    /**
     * A set of strings which are the fully-qualified names of all
     * known built-in immutable classes. Immutable.
     */
    private static final Set<String> builtinImmutables;

    static {
        Set<String> tmpImmutables = new HashSet<String>(12);
        tmpImmutables.add("java.lang.String");

        tmpImmutables.add("java.lang.Boolean");
        tmpImmutables.add("java.lang.Byte");
        tmpImmutables.add("java.lang.Character");
        tmpImmutables.add("java.lang.Short");
        tmpImmutables.add("java.lang.Integer");
        tmpImmutables.add("java.lang.Long");
        tmpImmutables.add("java.lang.Float");
        tmpImmutables.add("java.lang.Double");

        tmpImmutables.add("java.lang.Character.UnicodeBlock");
        tmpImmutables.add("java.lang.RuntimePermission");
        tmpImmutables.add("java.lang.StackTraceElement");

        builtinImmutables = Collections.unmodifiableSet(tmpImmutables);
    }

    /**
     * Visitor class that indicates whether a {@code TypeMirror}
     * represents an immutable type or not at a given strictness
     * level.
     * 
     * If the argument represents a primitive type, the result will be
     * {@code true}. If the argument represents a declared type, the
     * result will be {@code true} if it represents a class marked
     * immutable with a strictness equal to or higher than the given
     * level. If the argument represents a type parameter or an array,
     * the result will be {@code false}. If the argument is some other
     * type, an {@link UnsupportedOperationException} will be thrown.
     */
    private class IsImmutableTypeVisitor extends
            SimpleTypeVisitor6<Boolean, Immutable.Strictness> {

        /**
         * Throws an exception.
         * 
         * @param t
         *            the encountered type
         * @param p
         *            the strictness level required
         * @return nothing
         * @throws UnsupportedOperationException
         *             always
         */
        @Override
        protected Boolean defaultAction(TypeMirror t, Immutable.Strictness p) {
            throw new UnsupportedOperationException("encountered non-"
                    + "declared, non-primitive type " + t.toString());
        }

        /**
         * Prints an error message and returns {@code false}.
         * 
         * @param t
         *            the {@code ErrorType} encountered
         * @param p
         *            the strictness level required
         * @return {@code false}
         */
        @Override
        public Boolean visitError(ErrorType t, Immutable.Strictness p) {
            Messager messager = processingEnv.getMessager();
            messager.printMessage(Diagnostic.Kind.ERROR,
                    "Cannot process error type " + t.toString(), t.asElement());
            return false;
        }

        /**
         * Indicates whether the given {@code DeclaredType} is
         * immutable according to the given strictness level.
         * 
         * @param t
         *            the {@code DeclaredType} encountered
         * @param p
         *            the required strictness level
         * @return {@code true} if {@code t} is immutable by the by
         *         the requirements of {@code p}; {@code false}
         *         otherwise
         */
        @Override
        public Boolean visitDeclared(DeclaredType t, Immutable.Strictness p) {
            Types types = processingEnv.getTypeUtils();
            // TODO: is there any reason why we _can't_ make this
            // cast?
            TypeElement type = (TypeElement) t.asElement();
            Name name = type.getQualifiedName();
            Immutable annotation = type.getAnnotation(Immutable.class);

            if (builtinImmutables.contains(name.toString())) {
                return true;
            }
            if (Elements.isInterface(type, types)) {
                return false;
            }
            if (!Elements.isAnnotatedWith(type, Immutable.class)) {
                return false;
            }
            if (annotation.value().compareTo(p) < 0) {
                return false;
            }
            return true;
        }

        /**
         * Returns {@code true}.
         * 
         * @param t
         *            the {@code PrimitiveType} encountered
         * @param p
         *            the required strictness level
         * @return {@code true}
         */
        @Override
        public Boolean visitPrimitive(PrimitiveType t, Immutable.Strictness p) {
            return true;
        }

        /**
         * Returns {@code false}.
         * 
         * @param t
         *            the {@code TypeVariable} encountered
         * @param p
         *            the required strictness level
         * @return {@code false}
         */
        @Override
        public Boolean visitTypeVariable(TypeVariable t, Strictness p) {
            return false;
        }

        /**
         * Returns {@code false}.
         * 
         * @param t
         *            the {@code ArrayType} encountered
         * @param p
         *            the required strictness level
         * @return {@code false}
         */
        @Override
        public Boolean visitArray(ArrayType t, Strictness p) {
            return false;
        }
    }

    /**
     * Determines whether the given type is immutable or not.
     * 
     * @param type
     *            the type to check for immutability
     * @param requiredStrictness
     *            the strictness level required of {@code type}
     * @return {@code true} if {@code type} is immutable for
     *         {@code requiredStrictness} or higher; {@code false}
     *         otherwise
     */
    public boolean isImmutable(TypeMirror type,
            Immutable.Strictness requiredStrictness) {
        return new IsImmutableTypeVisitor().visit(type, requiredStrictness);
    }

    /**
     * Inspects an element which is marked immutable to determine
     * whether it meets the requirements. If an element does not meet
     * the requirements, an error message will be printed.
     * 
     * If this class tries to inspect an element that does not
     * represent a class, an error message will be printed.
     */
    private class ImmutableElementVisitor extends
            SimpleElementVisitor6<Void, Void> {

        /**
         * Inspects a field in a class marked immutable to determine
         * whether it meets the requirements. If an element does not
         * meet the requirements, an error message will be printed.
         * 
         * If this class tries to inspect an element that is not a
         * field or a method, an {@code UnsupportedOperationException}
         * will be thrown.
         */
        private class ImmutableClassFieldVisitor extends
                SimpleElementVisitor6<Void, Immutable> {

            /**
             * Throws an exception.
             * 
             * @param e
             *            the {@code Element} encountered
             * @param p
             *            is always {@code null}
             * @return nothing
             * @throws UnsupportedOperationException
             *             always
             */
            @Override
            protected Void defaultAction(Element e, Immutable p) {
                throw new UnsupportedOperationException(
                        "encountered non-executable, non-variable element "
                                + e.toString());
            }

            /**
             * Does nothing.
             * 
             * @param e
             *            the {@code ExecutableElement} encountered
             * @param p
             *            is always {@code null}
             * @return {@code null}
             */
            @Override
            public Void visitExecutable(ExecutableElement e, Immutable p) {
                return null;
            }

            /**
             * Checks the given {@code VariableElement} to ensure that
             * it meets the immutability requirements of its enclosing
             * class's immutable annotation.
             * 
             * @param e
             *            the {@code VariableElement} encountered
             * @param p
             *            is always {@code null}
             * @return {@code null}
             */
            @Override
            public Void visitVariable(VariableElement e, Immutable p) {
                Messager messager = processingEnv.getMessager();
                Set<Modifier> fieldModifiers = e.getModifiers();
                TypeMirror type = e.asType();
                Immutable annotation = p;
                Immutable.Strictness requiredStrictness =
                        annotation.value().getLowestRequiredStrictness();

                if (!annotation.value().arePublicMutableFieldsAllowed()
                        && Elements.isAnnotatedWith(e, Mutable.class)
                        && fieldModifiers.contains(Modifier.PUBLIC)) {
                    messager.printMessage(Diagnostic.Kind.ERROR,
                            "Public field is marked mutable.", e);
                }

                if (annotation.value().mustNonmutableFieldsBeFinal()
                        && !Elements.isAnnotatedWith(e, Mutable.class)
                        && !fieldModifiers.contains(Modifier.FINAL)) {
                    messager.printMessage(Diagnostic.Kind.ERROR,
                            "Nonmutable field is not final.", e);
                }

                if (!annotation.value().arePublicNonImmutableFieldsAllowed()
                        && fieldModifiers.contains(Modifier.PUBLIC)
                        && !Elements.isAnnotatedWith(e, Mutable.class)
                        && !isImmutable(type, requiredStrictness)) {
                    messager.printMessage(Diagnostic.Kind.ERROR,
                            "Public field is not immutable or has too low a "
                                    + "strictness level.", e);
                }
                if (!annotation.value().arePrivateNonImmutableFieldsAllowed()
                        && fieldModifiers.contains(Modifier.PRIVATE)
                        && !Elements.isAnnotatedWith(e, Mutable.class)
                        && !isImmutable(type, requiredStrictness)) {
                    messager.printMessage(Diagnostic.Kind.ERROR,
                            "Private field is not immutable or has too low a "
                                    + "strictness level.", e);
                }
                if (!annotation.value().areMutableFieldsAllowed()
                        && Elements.isAnnotatedWith(e, Mutable.class)) {
                    messager.printMessage(Diagnostic.Kind.ERROR,
                            "Field is marked mutable.", e);
                }

                return null;
            }
        }

        /**
         * Prints an error message.
         * 
         * @param e
         *            the {@code Element} encountered
         * @param p
         *            is always {@code null}
         * @return {@code null}
         */
        @Override
        protected Void defaultAction(Element e, Void p) {
            Messager messager = processingEnv.getMessager();
            messager.printMessage(Diagnostic.Kind.ERROR,
                    "Only classes may be marked immutable.", e);
            return null;
        }

        public void checkType(TypeElement e, Immutable i) {
            Types types = processingEnv.getTypeUtils();

            for (Element member : e.getEnclosedElements()) {
                new ImmutableClassFieldVisitor().visit(member, i);
            }

            if (Elements.hasConcreteSuperclass(e, types)) {
                TypeMirror sup = e.getSuperclass();
                checkType((TypeElement) types.asElement(sup), i);
            }
        }

        /**
         * Checks to ensure that the given {@code TypeElement} is not
         * an interface and otherwise meets the requirements of its
         * immutable annotation.
         * 
         * @param e
         *            the {@code TypeElement} encountered
         * @param p
         *            is always {@code null}
         * @return {@code null}
         */
        @Override
        public Void visitType(TypeElement e, Void p) {
            Types types = processingEnv.getTypeUtils();
            Messager messager = processingEnv.getMessager();
            Immutable annotation = e.getAnnotation(Immutable.class);
            if (annotation == null) {
                throw new UnsupportedOperationException(
                        "The given type is not marked immutable");
            }

            if (Elements.isInterface(e, types)) {
                messager.printMessage(Diagnostic.Kind.ERROR,
                        "An interface may not be marked immutable.", e);
            }

            if (annotation.value().mustClassBeFinal()
                    && !e.getModifiers().contains(Modifier.FINAL)) {
                messager.printMessage(Diagnostic.Kind.ERROR,
                        "Type is not declared final.", e);
            }

            checkType(e, annotation);

            return null;
        }
    }

    /**
     * Inspects an element marked mutable. This is only used to make
     * sure that it belongs to a class marked immutable. If it does
     * not, an error message is printed.
     * 
     * If this class tries to inspect an element that is not a class
     * or a field, an {@link UnsupportedOperationException} will be
     * thrown.
     */
    private class MutableElementVisitor extends
            SimpleElementVisitor6<Void, Void> {

        /**
         * Throws an exception.
         * 
         * @param e
         *            the {@code Element} encountered
         * @param p
         *            is always {@code null}
         * @return {@code null}
         * @throws UnsupportedOperationException
         *             always
         */
        @Override
        protected Void defaultAction(Element e, Void p) {
            throw new UnsupportedOperationException(
                    "encountered non-field element " + e.toString());
        }

        /**
         * Ensures that the given field belongs to an immutable class.
         * 
         * @param e
         *            the {@code VariableElement} encountered
         * @param p
         *            is always {@code null}
         * @return {@code null}
         */
        @Override
        public Void visitVariable(VariableElement e, Void p) {
            TypeElement type = (TypeElement) e.getEnclosingElement();
            if (!Elements.isAnnotatedWith(type, Immutable.class)) {
                Messager messager = processingEnv.getMessager();
                messager.printMessage(Diagnostic.Kind.ERROR,
                        "A mutable field must belong to an immutable class.", e);
            }

            return null;
        }
    }

    /**
     * Constructs a new {@code ImmutableProcessor}.
     */
    public ImmutableProcessor() {
        super();
    }

    /**
     * Processes a set of annotations. If any of the annotations in
     * the set are neither {@code Immutable} nor {@code Mutable}, an
     * exception will be thrown.
     * 
     * @param annotations
     *            the set of annotations to process
     * @param roundEnv
     *            the current processing {@code RoundEnvironment}
     * @return {@code true}
     */
    @Override
    public boolean process(Set<? extends TypeElement> annotations,
            RoundEnvironment roundEnv) {
        for (TypeElement annotation : annotations) {
            Set<? extends Element> elements;
            Name name = annotation.getQualifiedName();
            if (name.contentEquals(Immutable.class.getCanonicalName())) {
                elements = roundEnv.getElementsAnnotatedWith(Immutable.class);
                for (Element e : elements) {
                    new ImmutableElementVisitor().visit(e);
                }
            } else if (name.contentEquals(Mutable.class.getCanonicalName())) {
                elements = roundEnv.getElementsAnnotatedWith(Mutable.class);
                for (Element e : elements) {
                    new MutableElementVisitor().visit(e);
                }
            } else {
                throw new UnsupportedOperationException(
                        "encountered unknown annotation type "
                                + name.toString() + ".");
            }
        }
        return true;
    }
}
